#include "libmri.h"

#include <omp.h>

void plan_fft2c (fftwf_plan &forward, fftwf_plan &backward,
                 cplxf *X, int xs1, int xs2, int nch)
{
	fftwf_iodim xform[2] = {{xs1, 1, 1}, {xs2, xs1, xs1}};
	fftwf_iodim hmany[1] = {{nch, xs1*xs2, xs1*xs2}};

	forward = fftwf_plan_guru_dft (2, xform, 1, hmany,
	                               (fftwf_complex *) X, (fftwf_complex *) X,
	                               FFTW_FORWARD, FFTW_ESTIMATE);

	backward = fftwf_plan_guru_dft (2, xform, 1, hmany,
	                                (fftwf_complex *) X, (fftwf_complex *) X,
	                                FFTW_BACKWARD, FFTW_ESTIMATE);
}

void destroy_fft2c (fftwf_plan &forward, fftwf_plan &backward)
{
	fftwf_destroy_plan (forward);
	fftwf_destroy_plan (backward);
}

void fft2c_exec (fftwf_plan& direction, float scale,
                 cplxf *X, int xs1, int xs2, int nch)
{
	for (int ch = 0; ch < nch; ++ch)
		ifft2shift (&X[ch*xs1*xs2], xs1, xs2);

	fftwf_execute_dft (direction, (fftwf_complex*) X, (fftwf_complex*) X);

	for (int ch = 0; ch < nch; ++ch)
		fft2shift (&X[ch*xs1*xs2], xs1, xs2);

	for (int i = 0; i < xs1*xs2*nch; ++i)
		X[i] *= scale;
}

void ifft3c (cplxf *data, int *dsize)
{
	int ds0 = dsize[0], ds1 = dsize[1], ds2 = dsize[2], nch = dsize[3];

	#pragma omp parallel
	{
		#pragma omp for
		for (int ch = 0; ch < nch; ++ch)
			for (int i2 = 0; i2 < ds2; ++i2)
				for (int i1 = 0; i1 < ds1; ++i1)
					ifftshift (&data[ds0*(i1 + ds1*(i2 + ds2*ch))], ds0, 1);

		#pragma omp for
		for (int ch = 0; ch < nch; ++ch)
			for (int i2 = 0; i2 < ds2; ++i2)
				for (int i0 = 0; i0 < ds0; ++i0)
					ifftshift (&data[i0 + ds0*ds1*(i2 + ds2*ch)], ds1, ds0);

		#pragma omp for
		for (int ch = 0; ch < nch; ++ch)
			for (int i1 = 0; i1 < ds1; ++i1)
				for (int i0 = 0; i0 < ds0; ++i0)
					ifftshift (&data[i0 + ds0*(i1 + ds1*ds2*ch)], ds2, ds0*ds1);
	}

	fftwf_iodim xform[3] = {{ds0, 1, 1},
	                        {ds1, ds0, ds0},
	                        {ds2, ds0*ds1, ds0*ds1}};

	fftwf_iodim vectr[1] = {{nch, ds0*ds1*ds2, ds0*ds1*ds2}};

	fftwf_plan_with_nthreads (omp_get_max_threads());

	fftwf_plan inverse = fftwf_plan_guru_dft (3, xform, 1, vectr,
	                                          (fftwf_complex*) data,
	                                          (fftwf_complex*) data,
	                                          FFTW_BACKWARD, FFTW_ESTIMATE);
	fftwf_execute (inverse);
	fftwf_destroy_plan (inverse);

	#pragma omp parallel
	{
		#pragma omp for
		for (int ch = 0; ch < nch; ++ch)
			for (int i2 = 0; i2 < ds2; ++i2)
				for (int i1 = 0; i1 < ds1; ++i1)
					fftshift (&data[ds0*(i1 + ds1*(i2 + ds2*ch))], ds0, 1);

		#pragma omp for
		for (int ch = 0; ch < nch; ++ch)
			for (int i2 = 0; i2 < ds2; ++i2)
				for (int i0 = 0; i0 < ds0; ++i0)
					fftshift (&data[i0 + ds0*ds1*(i2 + ds2*ch)], ds1, ds0);

		#pragma omp for
		for (int ch = 0; ch < nch; ++ch)
			for (int i1 = 0; i1 < ds1; ++i1)
				for (int i0 = 0; i0 < ds0; ++i0)
					fftshift (&data[i0 + ds0*(i1 + ds1*ds2*ch)], ds2, ds0*ds1);
	}
}
